# Standard library imports
import json
import logging
import time

# Third-party imports
from flask import request

# Application-specific imports
# Importing configuration constants for notebook deployment
from myapp.config import NOTEBOOK_NAMESPACE, HUBSECRET, HOSTALIASES, JUPYTER_ACCOUNTS, IMAGE_PULL_POLICY, CRD_INFO
# Importing database models
from myapp.models.model_dataset import Dataset
from myapp.models.model_job import Repository, Images
# Importing Kubernetes client utility
from myapp.third.k8s.py_k8s import K8s
# Importing core utilities
from myapp.utils import core
# Importing device resource utility for GPU handling
from myapp.utils.device_resource import DeviceResource
# Importing environment utility to get service external IP
from myapp.utils.env import get_service_external_ip
# Importing exception handling utility
from myapp.utils.exception import ignore_exception


# Defines a function to reset or create a Theia/Jupyter notebook environment in Kubernetes.
# This function handles the deployment of a notebook pod, service, and an Istio VirtualService.
def reset_theia(notebook, session):
    # Generates a unique name for the notebook based on its representation.
    notebook_name = repr(notebook)

    # Initializes the Kubernetes client with the cluster configuration.
    k8s_client = K8s(notebook.cluster.get('KUBECONFIG', ''))
    # Sets the Kubernetes namespace for the notebook.
    namespace = NOTEBOOK_NAMESPACE
    # Initializes a list to hold external IP addresses for the service.
    service_external_ip = []
    # Retrieves the external IP of the service.
    external_ip = get_service_external_ip()
    # If an external IP is found, add it to the list.
    if external_ip:
        service_external_ip = [external_ip]

    # If no external IP is determined yet, check if the request host is an IP.
    if not service_external_ip:
        if core.checkip(request.host):
            service_external_ip = [request.host]
    # If still no external IP, raise an exception as it's required.
    if not service_external_ip:
        raise Exception('service_external_ip required')

    # Defines the default port for the notebook application.
    port = 3000
    # Extracts the username from the notebook's creator.
    username = notebook.created_by.username
    # Initializes command and working directory variables.
    command = None
    working_dir = None
    # Retrieves volume mount configurations from the notebook object.
    volume_mount = notebook.volume_mount
    # Defines a pre-command script to set up SSH and create mount directories.
    pre_command = (
        f'systemctl restart sshd; echo {notebook.ssh_secret} | passwd --stdin root; echo'
        f' root:{notebook.ssh_secret} | chpasswd; mkdir -p {notebook.mount};'
    )
    # Conditional logic to set up command and working directory based on IDE type.
    if notebook.ide_type in ['jupyter', 'bigdata', 'machinelearning', 'deeplearning']:
        # Sets the working directory to the notebook's mount point.
        working_dir = notebook.mount

        # Defines the command to start Jupyter Lab.
        command = [
            'sh',
            '-c',
            f'{pre_command} jupyter lab --notebook-dir={notebook.mount} --ip=0.0.0.0'
            f" --no-browser --allow-root --port={port} --NotebookApp.token=''"
            " --NotebookApp.password='' --ServerApp.disable_check_xsrf=True"
            " --NotebookApp.allow_origin='*' "
            # "--FileContentsManager.use_atomic_writing=False "
            # "--FileManagerMixin.use_atomic_writing=False "
            f'--NotebookApp.base_url={notebook.route_prefix}',
        ]

    # If the IDE type is Theia.
    elif notebook.ide_type == 'theia':
        # Defines the command to start Theia IDE.
        command = [
            'bash',
            '-c',
            (
                f'{pre_command} node /home/theia/src-gen/backend/main.js /home/project'
                f' --hostname=0.0.0.0 --port={port}'
            ),
        ]
        # Sets the working directory for Theia.
        working_dir = '/home/theia'
    # Logs the constructed command for debugging.
    logging.debug(f'command: {command}')
    # Logs the working directory for debugging.
    logging.debug(f'workingDir: {working_dir}')

    # Initializes image secrets with global HUBSECRET.
    image_secrets = HUBSECRET
    # Queries for user-specific hub secrets from the repository.
    user_hubsecrets = (
        session.query(Repository.hubsecret)
        .filter(Repository.created_by_fk == notebook.created_by.id)
        .all()
    )
    # Appends user-specific hub secrets to the image_secrets list if they are not already present.
    if user_hubsecrets:
        for hubsecret in user_hubsecrets:
            if hubsecret[0] not in image_secrets:
                image_secrets.append(hubsecret[0])

    # Defines Kubernetes pod labels for identification and management.
    labels = {'app': notebook_name, 'user': username, 'pod-type': 'notebook'}
    # 端口+0是jupyterlab，+1是sshd，+2，+3是预留的用户自己启动应用占用的端口
    # Defines environment variables for the notebook pod.
    env = {
        'NO_AUTH': 'true',  # Disables authentication for the notebook.
        'USERNAME': username,  # Sets the username environment variable.
        # Configures Node.js memory limit based on notebook's resource memory.
        'NODE_OPTIONS': '--max-old-space-size=%s'
                        % str(int(notebook.resource_memory.replace('G', '')) * 1024),
        'SSH_PORT': str(notebook.ssh_port),  # Sets the SSH port.
        # "PORT1": str(10000 + 10 * notebook.id + 2), # Reserved port 1
        # "PORT2": str(10000 + 10 * notebook.id + 3), # Reserved port 2
    }
    # If an external IP is available, add it to the environment variables.
    if service_external_ip:
        env['SERVICE_EXTERNAL_IP'] = service_external_ip[0]

    # Initializes lists for dataset IDs, datasets, and dataset paths.
    dataset_ids = []
    datasets = []
    datasets_path = []
    # Attempts to parse notebook expansion data, ignoring exceptions if it fails.
    with ignore_exception:
        expand = json.loads(notebook.expand)
        datasets = expand['datasets']

    # Extracts dataset IDs from the parsed datasets.
    for dataset in datasets:
        dataset_ids.append(dataset['id'])
    # If dataset IDs are found, query their paths from the database.
    if len(dataset_ids) > 0:
        res = session.query(Dataset.dataset_path).filter(Dataset.id.in_(dataset_ids)).all()
        # Appends retrieved dataset paths to the list.
        for r in res:
            datasets_path.append(r[0])
    # If dataset paths are found, join them into a comma-separated string and add to environment variables.
    if len(datasets_path) > 0:
        env['DATASETS_PATH'] = ','.join(datasets_path)

    # Initializes the image name from the notebook object.
    image_name = notebook.images
    # Queries the Images model to get the full repository tag for the image.
    img = session.query(Images).filter(Images.id == notebook.image_id).first()
    # If an image is found, update the image name.
    if img:
        image_name = img.repository_with_tag

    # Creates a debug pod in Kubernetes for the notebook.
    k8s_client.create_debug_pod(
        namespace=namespace,  # Kubernetes namespace.
        name=notebook_name,  # Name of the pod.
        labels=labels,  # Labels for the pod.
        command=command,  # Command to run inside the pod.
        args=None,  # Arguments for the command.
        volume_mount=volume_mount,  # Volume mounts for the pod.
        working_dir=working_dir,  # Working directory for the command.
        node_selector=notebook.get_node_selector(),  # Node selector for scheduling.
        resource_memory='0G~' + notebook.resource_memory,  # Memory resources.
        resource_cpu='0~' + notebook.resource_cpu,  # CPU resources.
        resource_gpu=DeviceResource.from_str(notebook.resource_gpu),  # GPU resources.
        image_pull_policy=IMAGE_PULL_POLICY,  # Image pull policy.
        image_pull_secrets=image_secrets,  # Image pull secrets.
        image=image_name,  # Container image name.
        hostAliases=HOSTALIASES,  # Host aliases for the pod.
        env=env,  # Environment variables for the container.
        privileged=None,  # Privileged mode setting.
        accounts=JUPYTER_ACCOUNTS,  # Jupyter accounts configuration.
        username=username,  # Username for the notebook.
        health=str(port),  # Health check port.
    )
    # Creates a Kubernetes service for the notebook pod.
    k8s_client.create_service(
        namespace=namespace,  # Kubernetes namespace.
        name=notebook_name,  # Name of the service.
        username=username,  # Username associated with the service.
        ports=[
            port,  # Port exposed by the service.
        ],
        selector=labels,  # Selector to match pods.
    )

    # Retrieves CRD (Custom Resource Definition) information for VirtualService.
    crd_info = CRD_INFO['virtualservice']
    # Constructs the CRD name for the VirtualService.
    crd_name = f"notebook-jupyter-{notebook_name.replace('_', '-')}"
    # Attempts to retrieve an existing VirtualService.
    vs_obj = k8s_client.get_one_crd(
        group=crd_info['group'],  # CRD group.
        version=crd_info['version'],  # CRD version.
        plural=crd_info['plural'],  # CRD plural name.
        namespace=namespace,  # Namespace of the CRD.
        name=crd_name,  # Name of the CRD instance.
    )
    # If an existing VirtualService is found, delete it.
    if vs_obj:
        k8s_client.delete_crd(
            group=crd_info['group'],
            version=crd_info['version'],
            plural=crd_info['plural'],
            namespace=namespace,
            name=crd_name,
        )
        # Waits for 1 second to ensure deletion propagates.
        time.sleep(1)

    # Constructs the JSON body for the new Istio VirtualService.
    crd_json = {
        'apiVersion': 'networking.istio.io/v1alpha3',  # API version for Istio networking.
        'kind': 'VirtualService',  # Kind of Kubernetes resource.
        'metadata': {'name': crd_name, 'namespace': namespace},  # Metadata including name and namespace.
        'spec': {
            'gateways': ['kubeflow/kubeflow-gateway'],  # Specifies the gateway to bind to.
            'hosts': ['*'],  # Hosts that the VirtualService applies to.
            'http': [
                {
                    'match': [{'uri': {'prefix': notebook.route_prefix}}],  # Matches requests based on URI prefix.
                    'rewrite': {'uri': notebook.route_prefix},  # Rewrites the URI.
                    'route': [
                        {
                            'destination': {
                                'host': f'{notebook_name}.{namespace}.svc.cluster.local',  # Destination service host.
                                'port': {'number': port},  # Destination service port.
                            }
                        }
                    ],
                    'timeout': '300s',  # Request timeout.
                }
            ],
        },
    }
    # Creates the VirtualService CRD in Kubernetes.
    crd = k8s_client.create_crd(
        group=crd_info['group'],  # CRD group.
        version=crd_info['version'],  # CRD version.
        plural=crd_info['plural'],  # CRD plural name.
        namespace=namespace,  # Namespace for the CRD.
        body=crd_json,  # JSON body of the CRD.
    )
    # 边缘模式时，需要根据项目组中的配置设置代理ip
    # If in edge mode and external IP is available, create an external service.
    if service_external_ip:
        # Constructs a name for the external service.
        service_external_name = (notebook_name + '-external').lower()[:60].strip('-')
        # Creates an external Kubernetes service.
        k8s_client.create_service(
            namespace=namespace,  # Kubernetes namespace.
            name=service_external_name,  # Name of the external service.
            username=username,  # Username associated with the service.
            ports=[
                [notebook.ssh_port, 22],  # SSH port mapping.
                [notebook.ssh_port + 1, 8080],  # Additional port for user debugging.
            ],  # ssh_port+1 端口用于用户调试
            selector=labels,  # Selector to match pods.
            external_ip=service_external_ip,  # External IP for the service.
        )
    # Returns the created VirtualService CRD object.
    return crd
